{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import gym\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ENV = \"CartPole-v0\"\n",
    "\n",
    "CAPCITY = 2000\n",
    "\n",
    "UPDAT_PRMT_PERIOD = 100\n",
    "REPLAY_PERIOD = 3\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "EPISODE = 10000\n",
    "MAX_STEP = 300\n",
    "\n",
    "ALPHA = 0.5\n",
    "BETA_INIT = 0.1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# proportional variant / used in class Memory\n",
    "class SumTree():\n",
    "    def __init__(self ,  ):\n",
    "        self.capcity = CAPCITY\n",
    "        self.alpha = ALPHA\n",
    "    \n",
    "        # struct of SumTree & memory for the transition\n",
    "        self.tree = np.zeros( 2 * self.capcity - 1)\n",
    "        self.data = [None] * self.capcity\n",
    "        \n",
    "        # pointer for the position\n",
    "        self.pointer = 0\n",
    "    \n",
    "    # add new priority in leaf_node\n",
    "    def add_leaf_node(self , p_alpha , transition):\n",
    "        leaf_idex = self.pointer + self.capcity - 1\n",
    "                \n",
    "        self.data[self.pointer] = transition\n",
    "        self.update_leaf_node(leaf_idex , p_alpha)\n",
    "    \n",
    "        self.pointer += 1\n",
    "        if self.pointer >= self.capcity: # ！not self.capcity-1 \n",
    "            self.pointer = 0\n",
    "    \n",
    "    # update leaf_node according leaf_idex\n",
    "    def update_leaf_node(self , leaf_idex , p_alpha):\n",
    "        change = p_alpha - self.tree[leaf_idex]\n",
    "        self.tree[leaf_idex] = p_alpha \n",
    "        self._update_parent_node(change , leaf_idex )\n",
    "            \n",
    "    # update the value of sum p in parent node\n",
    "    def _update_parent_node(self , change , child_idex ):\n",
    "        parent_idex = (child_idex - 1) // 2\n",
    "        \n",
    "        self.tree[parent_idex] += change \n",
    "        \n",
    "        if parent_idex != 0:\n",
    "            self._update_parent_node(change , parent_idex)    \n",
    "        \n",
    "    # sampling to get leaf idex and transition\n",
    "    def sampling(self , sample_idex):\n",
    "        leaf_idex = self._retrieve(sample_idex)\n",
    "        data_idex = leaf_idex - self.capcity + 1\n",
    "        \n",
    "        return [leaf_idex , self.tree[leaf_idex] , self.data[data_idex] ]\n",
    "        \n",
    "    # retrieve with O(log n)\n",
    "    def _retrieve(self , sample_idex , node_idex = 0):\n",
    "        left_child_idex = node_idex * 2 + 1\n",
    "        right_child_idex = left_child_idex + 1\n",
    "        \n",
    "        if left_child_idex >= len(self.tree):  # ! must be  >= \n",
    "            return node_idex\n",
    "        \n",
    "        if self.tree[left_child_idex] == self.tree[right_child_idex]: \n",
    "            return self._retrieve(sample_idex , np.random.choice([left_child_idex , right_child_idex]))\n",
    "        if self.tree[left_child_idex] > sample_idex:\n",
    "            return self._retrieve(sample_idex , node_idex = left_child_idex)\n",
    "        else:\n",
    "            return self._retrieve(sample_idex - self.tree[left_child_idex] , node_idex = right_child_idex )\n",
    "        \n",
    "    # sum of p in root node\n",
    "    def root_priority(self):\n",
    "        return self.tree[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Memory():\n",
    "    def __init__(self , tree_epsilon = 0.01):\n",
    "        self.epsilon = tree_epsilon\n",
    "        self.p_init = 1. \n",
    "        self.beta = BETA_INIT\n",
    "        self.beta_change_step = 0.001\n",
    "        self.capcity = CAPCITY\n",
    "        \n",
    "        self.sum_tree = SumTree() \n",
    "        \n",
    "    # store transition & priority before replay\n",
    "    def store(self , transition):\n",
    "        p_max = np.max(self.sum_tree.tree[-self.capcity:])\n",
    "        if p_max == 0:\n",
    "            p_max = self.p_init\n",
    "        self.sum_tree.add_leaf_node(p_max , transition)\n",
    "        \n",
    "    # update SumTree\n",
    "    def update(self , leaf_idex , td_error  ):\n",
    "        p = np.abs(td_error) + self.epsilon\n",
    "        p_alpha = np.power(p , ALPHA)\n",
    "        \n",
    "        for i in range(len(leaf_idex)):\n",
    "            self.sum_tree.update_leaf_node(leaf_idex[i] , p_alpha[i] )\n",
    "        \n",
    "    # sample\n",
    "    def sampling(self , batch_size ):\n",
    "        batch_idex = []\n",
    "        batch_transition = []\n",
    "        batch_ISweight = []\n",
    "        \n",
    "        segment = self.sum_tree.root_priority() / batch_size\n",
    "        \n",
    "        for i in range(batch_size):\n",
    "            low = segment * i\n",
    "            high = segment * (i + 1)\n",
    "            sample_idex = np.random.uniform(low , high)\n",
    "            idex , p_alpha , transition  = self.sum_tree.sampling(sample_idex)\n",
    "            prob = p_alpha / self.sum_tree.root_priority()\n",
    "            batch_ISweight.append( np.power(self.capcity * prob , -self.beta) )\n",
    "            batch_idex.append( idex )\n",
    "            batch_transition.append( transition)\n",
    "\n",
    "        i_maxiwi = np.power(self.capcity * np.min(self.sum_tree.tree[-self.capcity:]) / self.sum_tree.root_priority() , self.beta)\n",
    "\n",
    "        batch_ISweight = np.array(batch_ISweight) * i_maxiwi \n",
    "        \n",
    "        return batch_idex , batch_transition , batch_ISweight\n",
    "        \n",
    "    # change beta\n",
    "    def change_beta(self):\n",
    "        self.beta -= self.beta_change_step\n",
    "        return np.min(1 , self.beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class AgentModel():\n",
    "    def __init__(self , env , sess):\n",
    "        self.state_dim = env.observation_space.shape[0]\n",
    "        self.action_dim = env.action_space.n\n",
    "        self.step_size = 0.01\n",
    "        self.gamma = 0.5\n",
    "        self.sess = sess\n",
    "        self.epsilon = 0.8\n",
    "        self.build_net()\n",
    "        self.sess.run(tf.global_variables_initializer())\n",
    "        \n",
    "    # net_frame\n",
    "    def net_frame(self , clt_name , inputs ):\n",
    "        weight_init = tf.random_normal_initializer()\n",
    "        bias_init = tf.constant_initializer(0.1)\n",
    "        \n",
    "        layer1_units = 64\n",
    "        layer2_units = 32\n",
    "        \n",
    "        with tf.variable_scope(\"layer1\"):\n",
    "            weights1 = tf.get_variable(\"weight\", initializer = weight_init , collections = clt_name , \n",
    "                                      shape = [self.state_dim , layer1_units ])\n",
    "            bias1 = tf.get_variable(\"bias\" , initializer = bias_init , collections = clt_name , shape = [layer1_units] )\n",
    "            wx_b1 = tf.matmul(inputs , weights1) + bias1\n",
    "            h1 = tf.nn.relu(wx_b1)\n",
    "            \n",
    "        with tf.variable_scope(\"layer2\"):\n",
    "            weights2 = tf.get_variable(\"weight\" , initializer = weight_init , collections = clt_name , \n",
    "                                        shape = [layer1_units , layer2_units])\n",
    "            bias2 = tf.get_variable(\"bias\" , initializer = weight_init , collections = clt_name , \n",
    "                                    shape = [layer2_units])\n",
    "            wx_b2 = tf.matmul(h1 , weights2) + bias2\n",
    "            h2 = tf.nn.relu(wx_b2)\n",
    "            \n",
    "        with tf.variable_scope(\"layer3\"):\n",
    "            weights3 = tf.get_variable(\"weight\" , initializer = weight_init , collections = clt_name , \n",
    "                                        shape = [layer2_units , self.action_dim])\n",
    "            bias3 = tf.get_variable(\"bias\" , initializer = weight_init , collections = clt_name , \n",
    "                                    shape = [self.action_dim])\n",
    "            q_out = tf.matmul(h2 , weights3) + bias3\n",
    "        \n",
    "        return q_out       \n",
    "                        \n",
    "    # build net\n",
    "    def build_net(self):\n",
    "        with tf.variable_scope(\"q_net\"):\n",
    "            clt_name_q = [\"q_net_prmts\" , tf.GraphKeys.GLOBAL_VARIABLES]\n",
    "            self.inputs_q = tf.placeholder( dtype = tf.float32 , shape = [None , self.state_dim] , name = \"q_inputs\" )\n",
    "            self.q_value = self.net_frame(clt_name_q , self.inputs_q)\n",
    "        \n",
    "        with tf.variable_scope(\"target_net\"):\n",
    "            clt_name_target = [\"target_net_prmts\" , tf.GraphKeys.GLOBAL_VARIABLES]\n",
    "            self.inputs_target = tf.placeholder( dtype = tf.float32 , shape = [None , self.state_dim] , name = \"q_target\")\n",
    "            self.q_target = self.net_frame(clt_name_target , self.inputs_target)\n",
    "            \n",
    "        with tf.variable_scope(\"loss\"):\n",
    "            self.target = tf.placeholder( dtype = tf.float32 , shape=[ None , self.action_dim ] , name=\"target\" )\n",
    "            self.ISweight = tf.placeholder( dtype = tf.float32 , shape = [ None , self.action_dim] , name = \"wi\")\n",
    "            \n",
    "            self.td_error = self.target - self.q_value\n",
    "            self.loss = tf.reduce_sum(self.ISweight * tf.squared_difference(self.target , self.q_value ))  # ***\n",
    "            \n",
    "        with tf.variable_scope(\"train\"):\n",
    "            self.train_op = tf.train.RMSPropOptimizer(self.step_size).minimize(self.loss)\n",
    "            \n",
    "    # training\n",
    "    def training(self , state , next_state , reward , action , batch_ISweight ):\n",
    "        q_target , q_value = self.sess.run([self.q_target , self.q_value] , \n",
    "                                          feed_dict = {self.inputs_q : state , self.inputs_target : next_state })        \n",
    "        \n",
    "        target = reward + self.gamma * np.max(q_target)\n",
    "        reform_target = q_value.copy()\n",
    "        batch_idex = np.arange(BATCH_SIZE)\n",
    "        reform_target[batch_idex , action] = target\n",
    "        \n",
    "        batch_ISweight = np.stack([batch_ISweight , batch_ISweight] , axis = -1 )\n",
    "        \n",
    "        _ , td_error , loss = self.sess.run([self.train_op , self.td_error , self.loss] ,\n",
    "                                            feed_dict = {self.ISweight : batch_ISweight , \n",
    "                                                         self.inputs_q : state ,\n",
    "                                                         self.target : reform_target })\n",
    "        return td_error\n",
    "    \n",
    "    # update target_net prmt\n",
    "    def update_target_prmts(self):\n",
    "        q_value_prmts = tf.get_collection(\"q_net_prmts\")\n",
    "        q_target_prmts = tf.get_collection(\"target_net_prmts\")\n",
    "        self.sess.run( [tf.assign(v , t) for v , t in zip(q_value_prmts , q_target_prmts) ])\n",
    "        print(\"updating target network prmts...\")\n",
    "        \n",
    "    # chose action\n",
    "    def chose_action(self , current_state):\n",
    "        state = current_state[np.newaxis , :]\n",
    "        q_value = self.sess.run(self.q_value , feed_dict={self.inputs_q : state})\n",
    "        if np.random.random() > self.epsilon:\n",
    "            action = np.argmax(q_value)\n",
    "        else:\n",
    "            action = np.random.randint(0 , self.action_dim)\n",
    "        return action\n",
    "    \n",
    "    def greedy_action(self , current_state):\n",
    "        current_state = current_state[np.newaxis , :]  \n",
    "        q = self.sess.run(self.q_value , feed_dict={self.inputs_q : current_state} ) \n",
    "        action_greedy = np.argmax(q)\n",
    "        return action_greedy\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train( env , agent , memory ):\n",
    "    Transition = collections.namedtuple(\"Transition\" , [\"state\" , \"action\" , \"reward\" , \"next_state\" , \"done\"])\n",
    "    replay_iter = 0\n",
    "    update_target_iter = 0\n",
    "    for episode in range(EPISODE):\n",
    "        state = env.reset()\n",
    "        for step in range(MAX_STEP):\n",
    "            action = agent.chose_action(state)\n",
    "            next_state , reward , done , _ = env.step(action)\n",
    "            replay_iter += 1\n",
    "            \n",
    "            memory.store(Transition(state , action , reward , next_state , done))\n",
    "            \n",
    "            if replay_iter > CAPCITY and replay_iter % REPLAY_PERIOD == 0:\n",
    "                batch_idex , batch_transition , batch_ISweight = memory.sampling(BATCH_SIZE)\n",
    "                batch_state ,  batch_action , batch_reward , batch_next_state , batch_done = map(np.array , zip(*batch_transition))\n",
    "                td_error = agent.training(batch_state , batch_next_state , batch_reward , batch_action , batch_ISweight)\n",
    "\n",
    "                batch_iter = np.arange(BATCH_SIZE)\n",
    "                memory.update(batch_idex , td_error[batch_iter , batch_action])\n",
    "\n",
    "            if replay_iter > CAPCITY and replay_iter % UPDAT_PRMT_PERIOD == 0:\n",
    "                agent.update_target_prmts()\n",
    "            \n",
    "            if done:\n",
    "                if episode % 50 == 0:\n",
    "                    print(\"episode: %d ,  step: %d\" %( episode , step ))\n",
    "                break\n",
    "            state = next_state\n",
    "            \n",
    "                \n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    env = gym.make(ENV)\n",
    "    memory = Memory()\n",
    "    sess = tf.Session()\n",
    "    agent = AgentModel(env , sess)\n",
    "    train( env , agent , memory )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
